package cn.iocoder.yudao.module.liveprice.websocket;

import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import lombok.RequiredArgsConstructor;
import lombok.extern.slf4j.Slf4j;
import org.springframework.stereotype.Component;
import org.springframework.web.socket.CloseStatus;
import org.springframework.web.socket.TextMessage;
import org.springframework.web.socket.WebSocketSession;
import org.springframework.web.socket.handler.TextWebSocketHandler;

import java.util.HashMap;
import java.util.Map;

/**
 * 价格WebSocket处理器
 */
@Slf4j
@Component
@RequiredArgsConstructor
public class PriceWebSocketHandler extends TextWebSocketHandler {

    private final PriceWebSocketManager webSocketManager;
    private final ObjectMapper objectMapper;
    
    /**
     * 会话属性信息
     */
    private final Map<String, String> sessionTopics = new HashMap<>();

    @Override
    public void afterConnectionEstablished(WebSocketSession session) {
        log.info("建立WebSocket连接: {}", session.getId());
    }

    @Override
    protected void handleTextMessage(WebSocketSession session, TextMessage message) throws Exception {
        String payload = message.getPayload();
        log.info("收到WebSocket消息: {}", payload);

        JsonNode jsonNode = objectMapper.readTree(payload);
        String action = jsonNode.path("action").asText();

        switch (action) {
            case "subscribe":
                String topic = jsonNode.path("topic").asText();
                webSocketManager.addSession(topic, session);
                sessionTopics.put(session.getId(), topic);
                session.sendMessage(new TextMessage("{\"status\":\"subscribed\",\"topic\":\"" + topic + "\"}"));
                break;
            case "unsubscribe":
                String topicToUnsubscribe = sessionTopics.get(session.getId());
                if (topicToUnsubscribe != null) {
                    webSocketManager.removeSession(topicToUnsubscribe, session);
                    sessionTopics.remove(session.getId());
                    session.sendMessage(new TextMessage("{\"status\":\"unsubscribed\",\"topic\":\"" + topicToUnsubscribe + "\"}"));
                }
                break;
            default:
                session.sendMessage(new TextMessage("{\"error\":\"未知操作\"}"));
        }
    }

    @Override
    public void afterConnectionClosed(WebSocketSession session, CloseStatus status) {
        log.info("关闭WebSocket连接: {}, 状态: {}", session.getId(), status);
        String topic = sessionTopics.get(session.getId());
        if (topic != null) {
            webSocketManager.removeSession(topic, session);
            sessionTopics.remove(session.getId());
        }
    }
} 